package net.corda.node.services.messaging

import net.corda.core.identity.CordaX500Name
import net.corda.core.internal.ThreadBox
import net.corda.core.messaging.CordaRPCOps
import net.corda.core.messaging.MessageRecipients
import net.corda.core.messaging.SingleMessageRecipient
import net.corda.core.node.services.PartyInfo
import net.corda.core.serialization.SerializationDefaults
import net.corda.core.serialization.SingletonSerializeAsToken
import net.corda.core.serialization.deserialize
import net.corda.core.serialization.serialize
import net.corda.core.utilities.NetworkHostAndPort
import net.corda.core.utilities.contextLogger
import net.corda.core.utilities.sequence
import net.corda.core.utilities.trace
import net.corda.node.VersionInfo
import net.corda.node.services.config.NodeConfiguration
import net.corda.node.services.statemachine.StateMachineManagerImpl
import net.corda.node.utilities.AffinityExecutor
import net.corda.node.utilities.AppendOnlyPersistentMap
import net.corda.node.utilities.PersistentMap
import net.corda.nodeapi.internal.ArtemisMessagingComponent.*
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.P2P_QUEUE
import net.corda.nodeapi.internal.persistence.CordaPersistence
import net.corda.nodeapi.internal.persistence.NODE_DATABASE_PREFIX
import org.apache.activemq.artemis.api.core.ActiveMQObjectClosedException
import org.apache.activemq.artemis.api.core.Message.*
import org.apache.activemq.artemis.api.core.RoutingType
import org.apache.activemq.artemis.api.core.SimpleString
import org.apache.activemq.artemis.api.core.client.ClientConsumer
import org.apache.activemq.artemis.api.core.client.ClientMessage
import java.security.PublicKey
import java.time.Instant
import java.util.*
import java.util.concurrent.*
import javax.annotation.concurrent.ThreadSafe
import javax.persistence.Column
import javax.persistence.Entity
import javax.persistence.Id
import javax.persistence.Lob

// TODO: Stop the wallet explorer and other clients from using this class and get rid of persistentInbox

/**
 * This class implements the [MessagingService] API using Apache Artemis, the successor to their ActiveMQ product.
 * Artemis is a message queue broker and here we run a client connecting to the specified broker instance
 * [ArtemisMessagingServer]. It's primarily concerned with peer-to-peer messaging.
 *
 * Message handlers are run on the provided [AffinityExecutor] synchronously, that is, the Artemis callback threads
 * are blocked until the handler is scheduled and completed. This allows backpressure to propagate from the given
 * executor through into Artemis and from there, back through to senders.
 *
 * An implementation of [CordaRPCOps] can be provided. If given, clients using the CordaMQClient RPC library can
 * invoke methods on the provided implementation. There is more documentation on this in the docsite and the
 * CordaRPCClient class.
 *
 * @param serverAddress The address of the broker instance to connect to (might be running in the same process).
 * @param myIdentity The public key to be used as the ArtemisMQ address and queue name for the node.
 * @param nodeExecutor An executor to run received message tasks upon.
 * @param advertisedAddress The node address for inbound connections, advertised to the network map service and peers.
 * If not provided, will default to [serverAddress].
 */
@ThreadSafe
class P2PMessagingClient(config: NodeConfiguration,
                         private val versionInfo: VersionInfo,
                         serverAddress: NetworkHostAndPort,
                         private val myIdentity: PublicKey,
                         private val nodeExecutor: AffinityExecutor.ServiceAffinityExecutor,
                         private val database: CordaPersistence,
                         advertisedAddress: NetworkHostAndPort = serverAddress,
                         private val maxMessageSize: Int
) : SingletonSerializeAsToken(), MessagingService {
    companion object {
        private val log = contextLogger()
        // This is a "property" attached to an Artemis MQ message object, which contains our own notion of "topic".
        // We should probably try to unify our notion of "topic" (really, just a string that identifies an endpoint
        // that will handle messages, like a URL) with the terminology used by underlying MQ libraries, to avoid
        // confusion.
        private val topicProperty = SimpleString("platform-topic")
        private val sessionIdProperty = SimpleString("session-id")
        private val cordaVendorProperty = SimpleString("corda-vendor")
        private val releaseVersionProperty = SimpleString("release-version")
        private val platformVersionProperty = SimpleString("platform-version")
        private val amqDelayMillis = System.getProperty("amq.delivery.delay.ms", "0").toInt()
        private val messageMaxRetryCount: Int = 3

        fun createProcessedMessage(): AppendOnlyPersistentMap<UUID, Instant, ProcessedMessage, String> {
            return AppendOnlyPersistentMap(
                    toPersistentEntityKey = { it.toString() },
                    fromPersistentEntity = { Pair(UUID.fromString(it.uuid), it.insertionTime) },
                    toPersistentEntity = { key: UUID, value: Instant ->
                        ProcessedMessage().apply {
                            uuid = key.toString()
                            insertionTime = value
                        }
                    },
                    persistentEntityClass = ProcessedMessage::class.java
            )
        }

        fun createMessageToRedeliver(): PersistentMap<Long, Pair<Message, MessageRecipients>, RetryMessage, Long> {
            return PersistentMap(
                    toPersistentEntityKey = { it },
                    fromPersistentEntity = {
                        Pair(it.key,
                                Pair(it.message.deserialize(context = SerializationDefaults.STORAGE_CONTEXT),
                                        it.recipients.deserialize(context = SerializationDefaults.STORAGE_CONTEXT))
                        )
                    },
                    toPersistentEntity = { _key: Long, (_message: Message, _recipient: MessageRecipients): Pair<Message, MessageRecipients> ->
                        RetryMessage().apply {
                            key = _key
                            message = _message.serialize(context = SerializationDefaults.STORAGE_CONTEXT).bytes
                            recipients = _recipient.serialize(context = SerializationDefaults.STORAGE_CONTEXT).bytes
                        }
                    },
                    persistentEntityClass = RetryMessage::class.java
            )
        }

        private class NodeClientMessage(override val topicSession: TopicSession, override val data: ByteArray, override val uniqueMessageId: UUID) : Message {
            override val debugTimestamp: Instant = Instant.now()
            override fun toString() = "$topicSession#${String(data)}"
        }
    }

    private class InnerState {
        var running = false
        var p2pConsumer: ClientConsumer? = null
    }

    private val messagesToRedeliver = database.transaction {
        createMessageToRedeliver()
    }

    private val scheduledMessageRedeliveries = ConcurrentHashMap<Long, ScheduledFuture<*>>()

    /** A registration to handle messages of different types */
    data class Handler(val topicSession: TopicSession,
                       val callback: (ReceivedMessage, MessageHandlerRegistration) -> Unit) : MessageHandlerRegistration

    private val cordaVendor = SimpleString(versionInfo.vendor)
    private val releaseVersion = SimpleString(versionInfo.releaseVersion)
    /** An executor for sending messages */
    private val messagingExecutor = AffinityExecutor.ServiceAffinityExecutor("Messaging", 1)

    override val myAddress: SingleMessageRecipient = NodeAddress(myIdentity, advertisedAddress)
    private val messageRedeliveryDelaySeconds = config.messageRedeliveryDelaySeconds.toLong()
    private val artemis = ArtemisMessagingClient(config, serverAddress, maxMessageSize)
    private val state = ThreadBox(InnerState())
    private val handlers = CopyOnWriteArrayList<Handler>()

    private val processedMessages = createProcessedMessage()

    @Entity
    @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}message_ids")
    class ProcessedMessage(
            @Id
            @Column(name = "message_id", length = 36)
            var uuid: String = "",

            @Column(name = "insertion_time")
            var insertionTime: Instant = Instant.now()
    )

    @Entity
    @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}message_retry")
    class RetryMessage(
            @Id
            @Column(name = "message_id", length = 36)
            var key: Long = 0,

            @Lob
            @Column
            var message: ByteArray = ByteArray(0),

            @Lob
            @Column
            var recipients: ByteArray = ByteArray(0)
    )

    fun start() {
        state.locked {
            val session = artemis.start().session
            // Create a queue, consumer and producer for handling P2P network messages.
            p2pConsumer = session.createConsumer(P2P_QUEUE)
        }

        resumeMessageRedelivery()
    }

    private fun resumeMessageRedelivery() {
        messagesToRedeliver.forEach { retryId, (message, target) ->
            send(message, target, retryId)
        }
    }

    private val shutdownLatch = CountDownLatch(1)

    private fun processMessage(consumer: ClientConsumer): Boolean {
        // Two possibilities here:
        //
        // 1. We block waiting for a message and the consumer is closed in another thread. In this case
        //    receive returns null and we break out of the loop.
        // 2. We receive a message and process it, and stop() is called during delivery. In this case,
        //    calling receive will throw and we break out of the loop.
        //
        // It's safe to call into receive simultaneous with other threads calling send on a producer.
        val artemisMessage: ClientMessage = try {
            consumer.receive()
        } catch (e: ActiveMQObjectClosedException) {
            null
        } ?: return false

        val message: ReceivedMessage? = artemisToCordaMessage(artemisMessage)
        if (message != null)
            deliver(message)

        // Ack the message so it won't be redelivered. We should only really do this when there were no
        // transient failures. If we caught an exception in the handler, we could back off and retry delivery
        // a few times before giving up and redirecting the message to a dead-letter address for admin or
        // developer inspection. Artemis has the features to do this for us, we just need to enable them.
        //
        // TODO: Setup Artemis delayed redelivery and dead letter addresses.
        //
        // ACKing a message calls back into the session which isn't thread safe, so we have to ensure it
        // doesn't collide with a send here. Note that stop() could have been called whilst we were
        // processing a message but if so, it'll be parked waiting for us to count down the latch, so
        // the session itself is still around and we can still ack messages as a result.
        state.locked {
            artemisMessage.acknowledge()
        }
        return true
    }

    /**
     * Starts the p2p event loop: this method only returns once [stop] has been called.
     */
    fun run() {
        try {
            val consumer = state.locked {
                check(artemis.started != null) { "start must be called first" }
                check(!running) { "run can't be called twice" }
                running = true
                // If it's null, it means we already called stop, so return immediately.
                p2pConsumer ?: return
            }

            while (processMessage(consumer)) { }
        } finally {
            shutdownLatch.countDown()
        }
    }

    private fun artemisToCordaMessage(message: ClientMessage): ReceivedMessage? {
        try {
            val topic = message.required(topicProperty) { getStringProperty(it) }
            val sessionID = message.required(sessionIdProperty) { getLongProperty(it) }
            val user = requireNotNull(message.getStringProperty(HDR_VALIDATED_USER)) { "Message is not authenticated" }
            val platformVersion = message.required(platformVersionProperty) { getIntProperty(it) }
            // Use the magic deduplication property built into Artemis as our message identity too
            val uuid = message.required(HDR_DUPLICATE_DETECTION_ID) { UUID.fromString(message.getStringProperty(it)) }
            log.trace { "Received message from: ${message.address} user: $user topic: $topic sessionID: $sessionID uuid: $uuid" }

            return ArtemisReceivedMessage(TopicSession(topic, sessionID), CordaX500Name.parse(user), platformVersion, uuid, message)
        } catch (e: Exception) {
            log.error("Unable to process message, ignoring it: $message", e)
            return null
        }
    }

    private inline fun <T> ClientMessage.required(key: SimpleString, extractor: ClientMessage.(SimpleString) -> T): T {
        require(containsProperty(key)) { "Missing $key" }
        return extractor(key)
    }

    private class ArtemisReceivedMessage(override val topicSession: TopicSession,
                                         override val peer: CordaX500Name,
                                         override val platformVersion: Int,
                                         override val uniqueMessageId: UUID,
                                         private val message: ClientMessage) : ReceivedMessage {
        override val data: ByteArray by lazy { ByteArray(message.bodySize).apply { message.bodyBuffer.readBytes(this) } }
        override val debugTimestamp: Instant get() = Instant.ofEpochMilli(message.timestamp)
        override fun toString() = "${topicSession.topic}#${data.sequence()}"
    }

    private fun deliver(msg: ReceivedMessage): Boolean {
        state.checkNotLocked()
        // Because handlers is a COW list, the loop inside filter will operate on a snapshot. Handlers being added
        // or removed whilst the filter is executing will not affect anything.
        val deliverTo = handlers.filter { it.topicSession.isBlank() || it.topicSession == msg.topicSession }
        try {
            // This will perform a BLOCKING call onto the executor. Thus if the handlers are slow, we will
            // be slow, and Artemis can handle that case intelligently. We don't just invoke the handler
            // directly in order to ensure that we have the features of the AffinityExecutor class throughout
            // the bulk of the codebase and other non-messaging jobs can be scheduled onto the server executor
            // easily.
            //
            // Note that handlers may re-enter this class. We aren't holding any locks and methods like
            // start/run/stop have re-entrancy assertions at the top, so it is OK.
            nodeExecutor.fetchFrom {
                database.transaction {
                    if (msg.uniqueMessageId in processedMessages) {
                        log.trace { "Discard duplicate message ${msg.uniqueMessageId} for ${msg.topicSession}" }
                    } else {
                        if (deliverTo.isEmpty()) {
                            // TODO: Implement dead letter queue, and send it there.
                            log.warn("Received message ${msg.uniqueMessageId} for ${msg.topicSession} that doesn't have any registered handlers yet")
                        } else {
                            callHandlers(msg, deliverTo)
                        }
                        // TODO We will at some point need to decide a trimming policy for the id's
                        processedMessages[msg.uniqueMessageId] = Instant.now()
                    }
                }
            }
        } catch (e: Exception) {
            log.error("Caught exception whilst executing message handler for ${msg.topicSession}", e)
        }
        return true
    }

    private fun callHandlers(msg: ReceivedMessage, deliverTo: List<Handler>) {
        for (handler in deliverTo) {
            handler.callback(msg, handler)
        }
    }

    /**
     * Initiates shutdown: if called from a thread that isn't controlled by the executor passed to the constructor
     * then this will block until all in-flight messages have finished being handled and acknowledged. If called
     * from a thread that's a part of the [net.corda.node.utilities.AffinityExecutor] given to the constructor,
     * it returns immediately and shutdown is asynchronous.
     */
    fun stop() {
        val running = state.locked {
            // We allow stop() to be called without a run() in between, but it must have at least been started.
            check(artemis.started != null)
            val prevRunning = running
            running = false
            val c = p2pConsumer ?: throw IllegalStateException("stop can't be called twice")
            try {
                c.close()
            } catch (e: ActiveMQObjectClosedException) {
                // Ignore it: this can happen if the server has gone away before we do.
            }
            p2pConsumer = null
            prevRunning
        }
        if (running && !nodeExecutor.isOnThread) {
            // Wait for the main loop to notice the consumer has gone and finish up.
            shutdownLatch.await()
        }
        // Only first caller to gets running true to protect against double stop, which seems to happen in some integration tests.
        if (running) {
            state.locked {
                artemis.stop()
            }
        }
    }

    override fun send(message: Message, target: MessageRecipients, retryId: Long?, sequenceKey: Any, acknowledgementHandler: (() -> Unit)?) {
        // We have to perform sending on a different thread pool, since using the same pool for messaging and
        // fibers leads to Netty buffer memory leaks, caused by both Netty and Quasar fiddling with thread-locals.
        messagingExecutor.fetchFrom {
            state.locked {
                val mqAddress = getMQAddress(target)
                val artemis = artemis.started!!
                val artemisMessage = artemis.session.createMessage(true).apply {
                    putStringProperty(cordaVendorProperty, cordaVendor)
                    putStringProperty(releaseVersionProperty, releaseVersion)
                    putIntProperty(platformVersionProperty, versionInfo.platformVersion)
                    putStringProperty(topicProperty, SimpleString(message.topicSession.topic))
                    putLongProperty(sessionIdProperty, message.topicSession.sessionID)
                    writeBodyBufferBytes(message.data)
                    // Use the magic deduplication property built into Artemis as our message identity too
                    putStringProperty(HDR_DUPLICATE_DETECTION_ID, SimpleString(message.uniqueMessageId.toString()))

                    // For demo purposes - if set then add a delay to messages in order to demonstrate that the flows are doing as intended
                    if (amqDelayMillis > 0 && message.topicSession.topic == StateMachineManagerImpl.sessionTopic.topic) {
                        putLongProperty(HDR_SCHEDULED_DELIVERY_TIME, System.currentTimeMillis() + amqDelayMillis)
                    }
                }
                log.trace {
                    "Send to: $mqAddress topic: ${message.topicSession.topic} " +
                            "sessionID: ${message.topicSession.sessionID} uuid: ${message.uniqueMessageId}"
                }
                artemis.producer.send(mqAddress, artemisMessage)
                retryId?.let {
                    database.transaction {
                        messagesToRedeliver.computeIfAbsent(it, { Pair(message, target) })
                    }
                    scheduledMessageRedeliveries[it] = messagingExecutor.schedule({
                        sendWithRetry(0, mqAddress, artemisMessage, it)
                    }, messageRedeliveryDelaySeconds, TimeUnit.SECONDS)

                }
            }
        }
        acknowledgementHandler?.invoke()
    }

    override fun send(addressedMessages: List<MessagingService.AddressedMessage>, acknowledgementHandler: (() -> Unit)?) {
        for ((message, target, retryId, sequenceKey) in addressedMessages) {
            send(message, target, retryId, sequenceKey, null)
        }
        acknowledgementHandler?.invoke()
    }

    private fun sendWithRetry(retryCount: Int, address: String, message: ClientMessage, retryId: Long) {
        fun ClientMessage.randomiseDuplicateId() {
            putStringProperty(HDR_DUPLICATE_DETECTION_ID, SimpleString(UUID.randomUUID().toString()))
        }

        log.trace { "Attempting to retry #$retryCount message delivery for $retryId" }
        if (retryCount >= messageMaxRetryCount) {
            log.warn("Reached the maximum number of retries ($messageMaxRetryCount) for message $message redelivery to $address")
            scheduledMessageRedeliveries.remove(retryId)
            return
        }

        message.randomiseDuplicateId()

        state.locked {
            log.trace { "Retry #$retryCount sending message $message to $address for $retryId" }
            artemis.started!!.producer.send(address, message)
        }

        scheduledMessageRedeliveries[retryId] = messagingExecutor.schedule({
            sendWithRetry(retryCount + 1, address, message, retryId)
        }, messageRedeliveryDelaySeconds, TimeUnit.SECONDS)
    }

    override fun cancelRedelivery(retryId: Long) {
        database.transaction {
            messagesToRedeliver.remove(retryId)
        }
        scheduledMessageRedeliveries[retryId]?.let {
            log.trace { "Cancelling message redelivery for retry id $retryId" }
            if (!it.isDone) it.cancel(true)
            scheduledMessageRedeliveries.remove(retryId)
        }
    }

    private fun getMQAddress(target: MessageRecipients): String {
        return if (target == myAddress) {
            // If we are sending to ourselves then route the message directly to our P2P queue.
            P2P_QUEUE
        } else {
            // Otherwise we send the message to an internal queue for the target residing on our broker. It's then the
            // broker's job to route the message to the target's P2P queue.
            val internalTargetQueue = (target as? ArtemisAddress)?.queueName ?: throw IllegalArgumentException("Not an Artemis address")
            createQueueIfAbsent(internalTargetQueue)
            internalTargetQueue
        }
    }

    /** Attempts to create a durable queue on the broker which is bound to an address of the same name. */
    private fun createQueueIfAbsent(queueName: String) {
        state.alreadyLocked {
            val session = artemis.started!!.session
            val queueQuery = session.queueQuery(SimpleString(queueName))
            if (!queueQuery.isExists) {
                log.info("Create fresh queue $queueName bound on same address")
                session.createQueue(queueName, RoutingType.MULTICAST, queueName, true)
            }
        }
    }

    override fun addMessageHandler(topic: String,
                                   sessionID: Long,
                                   callback: (ReceivedMessage, MessageHandlerRegistration) -> Unit): MessageHandlerRegistration {
        return addMessageHandler(TopicSession(topic, sessionID), callback)
    }

    override fun addMessageHandler(topicSession: TopicSession,
                                   callback: (ReceivedMessage, MessageHandlerRegistration) -> Unit): MessageHandlerRegistration {
        require(!topicSession.isBlank()) { "Topic must not be blank, as the empty topic is a special case." }
        val handler = Handler(topicSession, callback)
        handlers.add(handler)
        return handler
    }

    override fun removeMessageHandler(registration: MessageHandlerRegistration) {
        handlers.remove(registration)
    }

    override fun createMessage(topicSession: TopicSession, data: ByteArray, uuid: UUID): Message {
        // TODO: We could write an object that proxies directly to an underlying MQ message here and avoid copying.
        return NodeClientMessage(topicSession, data, uuid)
    }

    // TODO Rethink PartyInfo idea and merging PeerAddress/ServiceAddress (the only difference is that Service address doesn't hold host and port)
    override fun getAddressOfParty(partyInfo: PartyInfo): MessageRecipients {
        return when (partyInfo) {
            is PartyInfo.SingleNode -> NodeAddress(partyInfo.party.owningKey, partyInfo.addresses.first())
            is PartyInfo.DistributedNode -> ServiceAddress(partyInfo.party.owningKey)
        }
    }
}
